import numpy as np
from PIL import Image
import os

label_path = "../ref_label.txt"
label_list = np.loadtxt(label_path, dtype=np.str)
label_dict = dict(label_list)

mapping_matrix = np.array([[2.83269438e+00, -5.68796925e-03, 1.73472348e-18],
                           [-9.28598668e-03, 2.94509748e+00, -1.73472348e-18],
                           [3.33378047e+02, -9.36243388e+01, 1.00000000e+00]])
color_path_list = ["Kinect_FE_1COLOR",
                   "Kinect_FE_2COLOR",
                   "Kinect_NU_1COLOR",
                   "Kinect_NU_2COLOR",
                   "Kinect_OC_1COLOR",
                   "Kinect_OC_2COLOR",
                   "Kinect_PS_1COLOR",
                   "Kinect_PS_2COLOR"]

color_path_root = "D:\\RGBD_Face_datasets\\Lock3DFace-share\\color"
depth_path_root = "D:\\RGBD_Face_datasets\\Lock3DFace-share\\depth"
new_color_path_root = "../new_dataset/color"
new_depth_path_root = "../new_dataset/depth"
new_image_path_root = "../new_dataset/image"

for index in range(200, 1000):
    if str(index) in label_dict.keys():
        new_index = label_dict[str(index)]
        if not os.path.exists(new_image_path_root + "/{:0>3}".format(new_index)):
            os.makedirs(new_color_path_root + "/{:0>3}".format(new_index))
            os.makedirs(new_depth_path_root + "/{:0>3}".format(new_index))
            os.makedirs(new_image_path_root + "/{:0>3}".format(new_index))

        cur_index = 0
        for color_path in color_path_list:
            color_path_file = color_path_root + "/{:0>3}_".format(index) + color_path
            if os.path.exists(color_path_file):
                cur_index += 1
                color_image = np.array(Image.open(color_path_file + "/01.jpg"))
                mapped_color_image = np.zeros((424, 512, 3), dtype='uint8')
                for i in range(424):
                    for j in range(512):
                        coordinate = np.dot(np.array([j, i, 1]), mapping_matrix)
                        color_x = coordinate[0]
                        color_y = coordinate[1]
                        if not (color_x == '-inf' or color_y == '-inf'):
                            color_x = int(round(float(color_x)))
                            color_y = int(round(float(color_y)))
                            if 0 <= color_x < 1920 and 0 <= color_y < 1080:
                                mapped_color_image[i][j][0] = int(color_image[color_y][color_x][0])
                                mapped_color_image[i][j][1] = int(color_image[color_y][color_x][1])
                                mapped_color_image[i][j][2] = int(color_image[color_y][color_x][2])

                new_color_image = Image.fromarray(mapped_color_image)
                new_color_path = new_color_path_root + "/{:0>3}".format(new_index) + "/{:0>3}.jpg".format(cur_index)
                new_color_image.save(new_color_path)

                depth_path_file = depth_path_root + "/{:0>3}_".format(index) + color_path.replace("COLOR", "DEPTH")
                depth_image = np.array(Image.open(depth_path_file + "/01.png"))
                np.clip(depth_image, 500, 1000)
                depth_image = 255 * (depth_image - 500) / 500

                new_depth_image = Image.fromarray(np.uint8(depth_image))
                new_depth_path = new_depth_path_root + "/{:0>3}".format(new_index) + "/{:0>3}.png".format(cur_index)
                new_depth_image.save(new_depth_path)
                new_image = np.zeros((424, 512, 4), dtype=np.uint8)
                new_image[:, :, 0] = mapped_color_image[:, :, 0]
                new_image[:, :, 1] = mapped_color_image[:, :, 1]
                new_image[:, :, 2] = mapped_color_image[:, :, 2]
                new_image[:, :, 3] = new_depth_image
                new_image_path = new_image_path_root + "/{:0>3}".format(new_index) + "/{:0>3}.npy".format(cur_index)
                np.save(new_image_path, new_image)
